import torch
import torch.nn as nn
import torch.nn.functional as nn_functional
from model import getModel
from load_data import dataloader
import glob
from setting import device, epoch_num, init, my_print, start
import cv2

net, optimizer, loss_func, scheduler = init(getModel)


def train():
    net.fc.train()
    for epoch in range(epoch_num):
        _lr = optimizer.state_dict()['param_groups'][0]['lr']
        sum_acc = 0
        sum_loss = 0
        for i, (x, y) in enumerate(dataloader):
            x = x.to(device)
            y = y.to(device)
            out = net(x)
            loss = loss_func(out, y)

            loss.backward()  # 根据loss计算模型的梯度
            optimizer.step()  # 根据梯度调整模型的参数
            optimizer.zero_grad()  # 梯度归0，进行下一轮计算

            acc = (out.argmax(dim=1) == y).sum().item()

            my_print(verbose=0, epoch=epoch + 1, step=i + 1, size=len(dataloader), loss=loss.item(),
                     accuracy=acc / len(y), _lr=_lr)

            sum_loss += loss
            sum_acc += acc

        my_print(verbose=1, epoch=epoch + 1, loss=sum_loss.item(), accuracy=sum_acc / len(dataloader.dataset),
                 _lr=_lr)
        torch.save(net, f'../model/net_epoch{epoch + 1}.pkl')

        # 每次epoch后对学习率进行更新
        scheduler.step()


@torch.no_grad()  # 不计算模型梯度
def test_single(_net, epoch) -> (int, torch.Tensor):
    net.eval()  # 进入测试模式

    acc = 0
    loss = 0
    for i, (x, y) in enumerate(dataloader):
        x = x.to(device)
        y = y.to(device)
        out = _net(x)
        loss += loss_func(out, y)
        pred = out.argmax(dim=1)
        acc += (pred == y).sum().item()

        for j in range(len(pred)):
            if pred[j] != y[j]:
                _x = (x[j].permute(1, 2, 0) * 255).int().to('cpu').numpy()
                cv2.imwrite(f'../img/bad_{epoch}_{i}_{j}.jpg', _x)
                print(j)

    return acc, loss


def test():
    models_name = glob.glob(f'../model/net_epoch*.pkl')
    for model_name in models_name:
        epoch = int(model_name[:-4].split('_')[1][5:])

        new_net = torch.load(model_name).to(device)
        acc, loss = test_single(new_net, epoch)

        my_print(verbose=2, epoch=epoch, loss=loss.item(), accuracy=acc / len(dataloader.dataset))


if __name__ == '__main__':
    start(train, test)
